import cv2
import numpy as np
import sys
import os
from ImageContrastEnhancement.histogram_equalization.contrast import ImageContraster


def main():
    image_path = 'image/'
    try:
        images = [x for x in os.listdir(image_path)]
    except FileNotFoundError:
        msg = "Sorry, the file " + image_path + "does not exist."
        print(msg)
    else:
        if len(images) is 0:
            msg = "Sorry, the folder does not exist image."
            print(msg)
            sys.exit()
        # 采用限制对比度自适应直方图均衡化CLAHE算法
        method = "CLAHE"   # HE  AHE  CLAHE  Bright
        clahe_blocks = 8
        clahe_threshold = 10.0
        # contraster
        icter = ImageContraster()
        for img in images:
            image_name = img.split('.')[0]
            # img = Image.open(image_path+img)
            img = cv2.imread(image_path + img)
            clahe_eq_img = icter.enhance_contrast(img, method=method, blocks=clahe_blocks, threshold=clahe_threshold)
            cv2.imwrite(image_path+str(image_name)+"CLAHE_result.png", clahe_eq_img)
        print('Finish Enhancement~')


if __name__ == "__main__":
    main()


